import torch
from torch import nn
from torchvision import datasets
from tqdm import tqdm

from main.data import dataloader
from main.repository.resnet import DQN, BasicBlock


def get_img(t_d):
    dataiter = iter(t_d)
    return next(dataiter)


def look_cifar10(_model: nn.Module):
    # 数据集
    cifar10 = dataloader.data_set(datasets.CIFAR10, batch_size=1, num_workers=1)
    # 测试集
    test_loader = cifar10(train=False, shuffle=True)

    # 关闭梯度
    with torch.no_grad():
        images, labels = get_img(test_loader)

    outputs = _model(images)
    _, predicted = torch.max(outputs, 1)
    dataloader.check_out(images[0], labels[0],
                         f'ai：我看这是{dataloader.labelToStr[predicted[0]]}')


def dqn_look_cifar10():
    model = DQN(in_channels=3, block=BasicBlock, num_actions=10)
    model.load_state('model/DQN.pth')
    look_cifar10(model)


def dqn_cifar10_accuracy():
    model = DQN(in_channels=3, block=BasicBlock, num_actions=10)
    model.load_state('model/DQN.pth')
    _d = model.to_gpu()

    model.eval()  # 设置模型为评估模式

    # 数据集
    cifar10 = dataloader.data_set(datasets.CIFAR10, batch_size=100, num_workers=4)
    # 测试集
    test_loader = cifar10(train=False, shuffle=True)

    correct = 0
    total = 0

    with torch.no_grad():
        for i, data in enumerate(tqdm(test_loader, desc=f'验证'), 0):
            inputs, labels = data[0].to(_d), data[1].to(_d)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = correct / total
    print(f'准确率为: {100 * accuracy:.2f}%')


if __name__ == '__main__':
    # dqn_look_cifar10()
    # 5轮 准确率58，是我手艺不行还是模型不行？
    dqn_cifar10_accuracy()
